package com.oliver.proxy.midjourney.application.service;

import com.oliver.proxy.midjourney.domain.aggregate.TaskAggregate;
import com.oliver.proxy.midjourney.domain.entity.InsertPoint;
import com.oliver.proxy.midjourney.domain.entity.Prod;
import com.oliver.proxy.midjourney.domain.entity.TaskList;
import com.oliver.proxy.midjourney.domain.repository.TaskAggregateRepository;
import com.oliver.proxy.midjourney.infrastructure.mapper.ProdMapper;
import com.oliver.proxy.midjourney.infrastructure.mapper.TaskListMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import javax.annotation.PostConstruct;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

@Slf4j
@Service
public class TaskProcessingService {
    private static final int BATCH_SIZE = 4;
    private final TaskAggregateRepository taskAggregateRepository;
    private final TaskListMapper taskListMapper;
    private final ProdMapper prodMapper;

    public TaskProcessingService(TaskAggregateRepository taskAggregateRepository,
                               TaskListMapper taskListMapper,
                               ProdMapper prodMapper) {
        this.taskAggregateRepository = taskAggregateRepository;
        this.taskListMapper = taskListMapper;
        this.prodMapper = prodMapper;
    }

    private int calculateRequiredQuantity(int originalQty) {
        return (originalQty + BATCH_SIZE - 1) / BATCH_SIZE * BATCH_SIZE;
    }

    @PostConstruct
    @Scheduled(fixedRate = 5000)
    @Transactional
    public void processUnhandledTasks() {
        try {
            List<TaskList> unhandledTasks = taskListMapper.selectByStatus(0);
            for (TaskList task : unhandledTasks) {
                processTask(task);
            }
        } catch (Exception e) {
            log.error("Error processing unhandled tasks", e);
        }
    }

    private void processTask(TaskList task) {
        try {
            TaskAggregate taskAggregate = taskAggregateRepository.findById(task.getTaskId());
            if (taskAggregate == null) {
                return;
            }

            if (taskAggregate.isTypeZero()) {
                processTypeZeroTask(taskAggregate);
            } else if (taskAggregate.isTypeOne()) {
                processTypeOneTask(taskAggregate);
            }

            task.setTaskStatus(1);
            taskListMapper.update(task);
        } catch (Exception e) {
            log.error("Error processing task: {}", task.getTaskId(), e);
        }
    }

    private void processTypeZeroTask(TaskAggregate taskAggregate) {
        TaskList task = taskAggregate.getTaskList();
        String promptTemplate = task.getTaskPrompt();
        
        // 向上取整到4的倍数
        int requiredQty = calculateRequiredQuantity(task.getTaskGenQty());
        
        Map<String, String> insertPointMap = taskAggregate.getInsertPoints().stream()
                .collect(Collectors.toMap(InsertPoint::getPoint, InsertPoint::getContent));

        Set<Integer> uniquePlaceholders = extractUniquePlaceholders(promptTemplate);
        
        // 验证插入点
        for (Integer index : uniquePlaceholders) {
            String point = "{" + index + "}";
            if (!insertPointMap.containsKey(point)) {
                throw new IllegalArgumentException("Missing insert point for placeholder: " + point);
            }
        }

        List<List<String>> combinations = generateCombinations(insertPointMap, uniquePlaceholders);
        List<Prod> prods = new ArrayList<>();
        int currentIndex = 0;
        int totalBatches = requiredQty / BATCH_SIZE;

        // 按batch生成记录
        for (int batchNum = 1; batchNum <= totalBatches; batchNum++) {
            // 如果已经用完所有组合，重新开始
            if (currentIndex >= combinations.size()) {
                currentIndex = 0;
            }

            // 为当前batch生成prompt
            String batchPrompt = generatePrompt(promptTemplate, uniquePlaceholders, combinations.get(currentIndex));
            currentIndex++;

            // 为当前batch生成4条记录
            for (int i = 0; i < BATCH_SIZE; i++) {
                Prod prod = new Prod();
                prod.setTaskId(task.getTaskId());
                prod.setPrompt(batchPrompt);
                prod.setBatch(batchNum);
                prods.add(prod);
            }
        }

        if (!prods.isEmpty()) {
            prodMapper.batchInsert(prods);
        }
    }

    private String generatePrompt(String promptTemplate, Set<Integer> uniquePlaceholders, List<String> combination) {
        String finalPrompt = promptTemplate;
        Map<Integer, String> indexValueMap = new HashMap<>();
        
        for (int i = 0; i < combination.size(); i++) {
            indexValueMap.put(new ArrayList<>(uniquePlaceholders).get(i), combination.get(i));
        }
        
        for (Integer index : uniquePlaceholders) {
            String value = indexValueMap.get(index);
            String placeholder = "{" + index + "}";
            finalPrompt = finalPrompt.replace(placeholder, value);
        }
        
        return finalPrompt;
    }

    private void processTypeOneTask(TaskAggregate taskAggregate) {
        TaskList task = taskAggregate.getTaskList();
        
        // 向上取整到4的倍数
        int requiredQty = calculateRequiredQuantity(task.getTaskGenQty());
        List<Prod> prods = new ArrayList<>();
        
        // 按batch生成记录
        int totalBatches = requiredQty / BATCH_SIZE;
        for (int batchNum = 1; batchNum <= totalBatches; batchNum++) {
            for (int i = 0; i < BATCH_SIZE; i++) {
                Prod prod = new Prod();
                prod.setTaskId(task.getTaskId());
                prod.setBatch(batchNum);
                prods.add(prod);
            }
        }
        
        if (!prods.isEmpty()) {
            prodMapper.batchInsert(prods);
        }
    }

    private Set<Integer> extractUniquePlaceholders(String prompt) {
        Pattern pattern = Pattern.compile("\\{(\\d+)\\}");
        Matcher matcher = pattern.matcher(prompt);
        Set<Integer> placeholders = new TreeSet<>();
        while (matcher.find()) {
            placeholders.add(Integer.parseInt(matcher.group(1)));
        }
        return placeholders;
    }

    private List<List<String>> generateCombinations(Map<String, String> insertPointMap, Set<Integer> uniquePlaceholders) {
        List<List<String>> result = new ArrayList<>();
        result.add(new ArrayList<>());

        for (Integer index : uniquePlaceholders) {
            String point = "{" + index + "}";
            String content = insertPointMap.get(point);
            if (content == null) {
                continue;
            }

            String[] options = content.split("\\^");
            List<List<String>> newCombinations = new ArrayList<>();

            for (List<String> combination : result) {
                for (String option : options) {
                    List<String> newCombination = new ArrayList<>(combination);
                    newCombination.add(option.trim());
                    newCombinations.add(newCombination);
                }
            }

            result = newCombinations;
        }

        return result;
    }
}